import sys
import os
import argparse
import time
import zipfile
import shutil
import tempfile
import subprocess
import json
from pathlib import Path
from typing import Optional

# Import the new multi-website management components
from models import WebsiteProfile, create_default_website_profile, AuthMethod
from session_manager import SessionManager
from website_manager import WebsiteManager
from auth_manager import AuthenticationManager

# 设置Playwright浏览器路径
def setup_playwright_path():
    """
    配置Playwright使用系统Chrome浏览器
    """
    if getattr(sys, 'frozen', False):
        # 打包后的执行文件 - 优先使用系统Chrome
        print("打包版本：尝试使用系统Chrome浏览器")
        # 不设置PLAYWRIGHT_BROWSERS_PATH，让Playwright使用系统Chrome
    else:
        # 开发环境，使用默认路径
        print("开发环境：使用系统默认Playwright浏览器路径")

setup_playwright_path()
 
from config import AUTH_FILE_NAME, LOGIN_TYPE, USERNAME, PASSWORD, JSON_PATH, TARGET_URL, upload_file_to_oss

# Try to import Playwright with error handling
try:
    from playwright.sync_api import sync_playwright
    PLAYWRIGHT_AVAILABLE = True
except ImportError as e:
    print(f"Warning: Playwright not available: {e}")
    print("Some features may be limited. Install with: pip install playwright")
    PLAYWRIGHT_AVAILABLE = False
    sync_playwright = None

def get_system_chrome_profile_dir():
    """
    获取系统Chrome默认配置目录
    """
    try:
        if sys.platform == "win32":
            profile_paths = [
                os.path.expanduser(r"~\AppData\Local\Google\Chrome\User Data\Default"),
                os.path.expanduser(r"~\AppData\Local\Google\Chrome\User Data\Profile 1"),
            ]
        elif sys.platform == "darwin":  # macOS
            profile_paths = [
                os.path.expanduser("~/Library/Application Support/Google/Chrome/Default"),
                os.path.expanduser("~/Library/Application Support/Google/Chrome/Profile 1"),
            ]
        else:  # Linux
            profile_paths = [
                os.path.expanduser("~/.config/google-chrome/Default"),
                os.path.expanduser("~/.config/google-chrome/Profile 1"),
            ]
        
        for path in profile_paths:
            if os.path.exists(path):
                print(f"🎯 找到系统Chrome配置: {path}")
                return path
        
        print("⚠️  未找到系统Chrome配置目录")
        return None
        
    except Exception as e:
        print(f"❌ 获取系统Chrome配置目录失败: {e}")
        return None

def backup_extensions_data(profile_dir, username):
    """
    备份Chrome扩展程序数据
    """
    try:
        print("📦 正在备份扩展程序数据...")
        
        # 创建临时目录用于存放要备份的数据
        temp_backup_dir = f"temp_extensions_backup_{username}"
        os.makedirs(temp_backup_dir, exist_ok=True)
        
        try:
            # 要备份的扩展程序目录
            extension_dirs = [
                "Extensions",
                "Extension State", 
                "Extension Rules",
                "Local Extension Settings",
                "Sync Extension Settings",
                "Local Storage",
                "Session Storage", 
                "IndexedDB"
            ]
            
            # 要备份的扩展程序文件
            extension_files = [
                "Extension Cookies",
                "Extension Cookies-journal",
                "Web Data",
                "Web Data-journal"
            ]
            
            copied_count = 0
            
            # 处理目录
            for item in extension_dirs:
                source_path = os.path.join(profile_dir, item)
                target_path = os.path.join(temp_backup_dir, item)
                
                if os.path.exists(source_path) and os.path.isdir(source_path):
                    try:
                        shutil.copytree(source_path, target_path)
                        copied_count += 1
                        print(f"📁 已备份目录: {item}")
                    except Exception as e:
                        print(f"⚠️  备份目录 {item} 失败: {e}")
            
            # 处理文件
            for item in extension_files:
                source_path = os.path.join(profile_dir, item)
                target_path = os.path.join(temp_backup_dir, item)
                
                if os.path.exists(source_path) and os.path.isfile(source_path):
                    try:
                        # 确保目标目录存在
                        os.makedirs(os.path.dirname(target_path), exist_ok=True)
                        shutil.copy2(source_path, target_path)
                        copied_count += 1
                        print(f"📄 已备份文件: {item}")
                    except Exception as e:
                        print(f"⚠️  备份文件 {item} 失败: {e}")
            
            if copied_count == 0:
                print("⚠️  没有找到任何扩展程序数据")
                return False
            
            # 创建压缩包
            backup_filename = f"{username}_extensions_backup.zip"
            
            print(f"🗜️  正在创建压缩包...")
            with zipfile.ZipFile(backup_filename, 'w', zipfile.ZIP_DEFLATED) as zipf:
                for root, dirs, files in os.walk(temp_backup_dir):
                    for file in files:
                        file_path = os.path.join(root, file)
                        arcname = os.path.relpath(file_path, temp_backup_dir)
                        zipf.write(file_path, arcname)
            
            # 获取压缩包大小
            backup_size = os.path.getsize(backup_filename) / (1024 * 1024)  # MB
            print(f"📦 压缩包大小: {backup_size:.2f} MB")
            
            # 上传到云端
            try:
                print(f"☁️  正在上传扩展数据到云端...")
                upload_file_to_oss(backup_filename)
                print(f"✅ 扩展程序数据备份成功!")
                
                # 清理本地临时文件
                os.remove(backup_filename)
                return True
                
            except Exception as e:
                print(f"⚠️  上传到云端失败: {e}")
                print(f"💾 本地备份文件: {backup_filename}")
                return False
                
        finally:
            # 清理临时目录
            if os.path.exists(temp_backup_dir):
                shutil.rmtree(temp_backup_dir)
                
    except Exception as e:
        print(f"❌ 备份扩展程序数据失败: {e}")
        return False

def find_chrome_executable():
    """查找系统中的Chrome可执行文件"""
    possible_paths = [
        r"C:\Program Files\Google\Chrome\Application\chrome.exe",
        r"C:\Program Files (x86)\Google\Chrome\Application\chrome.exe",
        os.path.expanduser(r"~\AppData\Local\Google\Chrome\Application\chrome.exe"),
    ]
    
    for path in possible_paths:
        if os.path.isfile(path):
            print(f"✅ 找到Chrome: {path}")
            return path
    
    # 如果都没找到，尝试命令行方式
    try:
        result = subprocess.run(["chrome", "--version"], 
                               capture_output=True, text=True, timeout=5)
        if result.returncode == 0:
            print(f"✅ 找到Chrome: chrome")
            return "chrome"
    except Exception:
        pass
        
    print("❌ 未找到Chrome浏览器")
    return None

def multi_website_login(username: str, website_id: str, password: Optional[str] = None) -> bool:
    """
    Enhanced login function that supports multi-website management.
    
    Args:
        username: User identifier
        website_id: Website profile identifier  
        password: Optional password for form-based authentication
        
    Returns:
        True if login successful, False otherwise
    """
    try:
        # Initialize managers
        session_manager = SessionManager()
        website_manager = WebsiteManager()
        auth_manager = AuthenticationManager()
        
        # Get website profile
        profile = website_manager.get_profile(website_id)
        if not profile:
            print(f"❌ Website profile not found: {website_id}")
            print("Available profiles:")
            profiles = website_manager.list_profiles(include_templates=True)
            for p in profiles:
                print(f"   - {p.website_id}: {p.display_name}")
            return False
        
        print(f"🌐 Logging into: {profile.display_name}")
        print(f"📍 URL: {profile.target_url}")
        print(f"🔐 Auth Method: {profile.auth_method.value}")
        
        # Perform authentication
        session = auth_manager.authenticate_website(
            profile=profile,
            username=username,
            password=password,
            backup_extensions=True
        )
        
        if not session:
            print("❌ Authentication failed")
            return False
        
        # Store session in session manager
        stored_session = session_manager.create_session(
            website_id=website_id,
            username=username,
            auth_state=session.auth_state,
            extensions_data=session.extensions_data
        )
        
        if stored_session:
            print("✅ Login successful and session stored!")
            print(f"📝 Session ID: {stored_session.session_id}")
            print(f"⏰ Expires: {stored_session.expires_at.strftime('%Y-%m-%d %H:%M:%S') if stored_session.expires_at else 'Never'}")
            return True
        else:
            print("⚠️  Login successful but failed to store session")
            return False
            
    except Exception as e:
        print(f"❌ Multi-website login failed: {e}")
        return False


def legacy_login_with_extensions(username=None, password=None, target_url=None, backup_extensions=True):
    """
    使用系统Chrome浏览器登录，直接使用Playwright启动保证状态保存
    """
    from config import LOGIN_TYPE, TARGET_URL, upload_file_to_oss
    
    # 使用传入的参数或配置文件中的值
    if username is None:
        from config import USERNAME
        username = USERNAME
        
    if target_url is None:
        target_url = TARGET_URL
        
    auth_file_name = f"{username}_auth_state.json"
    
    try:
        # 查找系统Chrome可执行文件（像open.py一样）
        chrome_path = find_chrome_executable()
        if not chrome_path:
            print("❌ 错误：未找到Chrome浏览器")
            return False
        
        # 获取系统Chrome配置目录
        system_profile = get_system_chrome_profile_dir()
        profile_dir = system_profile if (system_profile and backup_extensions) else None
        
        if not profile_dir:
            timestamp = int(time.time())
            profile_dir = os.path.join(os.getcwd(), "temp_login_profile", f"profile_{timestamp}")
            os.makedirs(profile_dir, exist_ok=True)
        
        # 构建Chrome启动参数（像open.py一样）
        chrome_args = [
            chrome_path,
            f"--user-data-dir={profile_dir}",
            "--start-maximized",
            target_url
        ]
        
        print("🚀 正在启动系统自带的Chrome浏览器...")
        
        # 直接启动系统Chrome进程（像open.py一样）
        process = subprocess.Popen(
            chrome_args,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            stdin=subprocess.DEVNULL,
            creationflags=subprocess.CREATE_NEW_PROCESS_GROUP if sys.platform == 'win32' else 0
        )
        
        print(f"✅ 系统Chrome进程启动成功 (PID: {process.pid})")
        time.sleep(3)
        
        print("\n" + "="*60)
        print("📋 请完成以下操作：")
        print("1. 🔐 完成网站登录")
        print("2. 🔌 安装和配置需要的浏览器插件")
        print("3. 🔑 完成插件的登录和配置")
        print("4. ✅ 确认所有插件都已正常工作")
        print("5. ⌨️  完成后按回车键保存状态")
        print("="*60)
        
        input("\n按回车键保存登录状态...")
        
        # 关闭Chrome进程，然后使用Playwright连接相同配置目录来保存状态
        print("📌 请先关闭Chrome浏览器，然后按回车键继续保存登录状态...")
        input()
        
        try:
            process.wait(timeout=5)
        except subprocess.TimeoutExpired:
            process.terminate()
            time.sleep(2)
        
        # 使用Playwright连接相同的用户数据目录来保存状态
        try:
            with sync_playwright() as p:
                print("💾 正在使用Playwright保存登录状态...")
                
                context = p.chromium.launch_persistent_context(
                    user_data_dir=profile_dir,
                    headless=True,
                    channel="chrome",
                    args=['--disable-web-security', '--no-first-run']
                )
                
                page = context.new_page()
                page.goto(target_url, timeout=30000)
                time.sleep(3)
                
                context.storage_state(path=auth_file_name)
                context.close()
                
                # 验证保存状态
                if os.path.exists(auth_file_name):
                    with open(auth_file_name, 'r', encoding='utf-8') as f:
                        saved_state = json.load(f)
                        cookies_count = len(saved_state.get('cookies', []))
                        if cookies_count > 0:
                            print(f"✅ 登录状态保存成功！Cookies: {cookies_count} 个")
                        else:
                            print("⚠️  警告：登录状态为空")
                
        except Exception as save_error:
            print(f"⚠️  保存失败: {save_error}")
            return False
        
        # 上传和备份
        try:
            upload_file_to_oss(auth_file_name)
            print("✅ 上传成功")
        except Exception as e:
            print(f"⚠️  上传失败: {e}")
        
        if system_profile and backup_extensions:
            extensions_backup_success = backup_extensions_data(system_profile, username)
            if extensions_backup_success:
                print("🎉 扩展程序数据备份成功!")
        
        return True
            
    except Exception as e:
        print(f"❌ 登录失败: {e}")
        print("可能的解决方案:")
        print("1. 检查网络连接")
        print("2. 检查Chrome浏览器是否正常安装")
        print("3. 确保目标网站可以正常访问")
        print("4. 尝试以管理员身份运行")
        return False

def parse_arguments():
    """
    Parse command line arguments for both legacy and multi-website login.
    """
    parser = argparse.ArgumentParser(
        description='Multi-Website Login Tool - Enhanced Login Program',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''Usage Examples:
  # Legacy single-website login (backward compatibility)
  python login.py                                    # Use config.py settings
  python login.py -u 13800138000                    # Specify username
  python login.py -u user --url https://example.com # Specify URL
  
  # New multi-website login
  python login.py --website geekbi -u 13800138000   # Login to GeekBI profile
  python login.py --website tabcut -u user          # Login to TabCut profile
  python login.py --list-websites                   # Show available website profiles
        '''
    )
    
    # Multi-website arguments
    parser.add_argument('--website', '--site', dest='website_id',
                       help='Website profile ID for multi-website login')
    parser.add_argument('--list-websites', action='store_true',
                       help='List available website profiles')
    
    # User credentials
    parser.add_argument('-u', '--username', 
                       help='Login username (phone number or email)')
    parser.add_argument('-p', '--password', 
                       help='Login password (optional, will prompt if needed)')
    
    # URL specification (legacy compatibility)
    parser.add_argument('--url', '--target-url', dest='target_url',
                       help='Target website URL (legacy mode)')
    
    # Options
    parser.add_argument('--no-extensions', action='store_true',
                       help='Disable browser extensions backup')
    
    return parser.parse_args()

if __name__ == "__main__":
    try:
        # 解析命令行参数
        args = parse_arguments()
        
        print("=" * 70)
        print("多点登录工具")
        print("=" * 70)
        print("🎯 特性: 支持网站登录 + 状态同步")
        print("=" * 70)
        
        # 获取用户名、密码和目标网址
        username = args.username
        password = args.password
        target_url = args.target_url
        backup_extensions = not args.no_extensions
        
        # 如果没有提供参数，则使用配置文件中的值
        if username is None:
            from config import USERNAME
            username = USERNAME
            
        if target_url is None:
            try:
                from config import TARGET_URL
                target_url = TARGET_URL
            except ImportError:
                target_url = "https://www.geekbi.com"  # 默认值
            
        # 密码为可选参数，如果未提供且配置文件中有密码，则使用配置文件中的密码
        if password is None:
            try:
                from config import PASSWORD
                password = PASSWORD if PASSWORD else None
            except (ImportError, AttributeError):
                password = None
            
        print(f"👤 用户名: {username}")
        if args.username:
            print("   (来源: 命令行参数)")
        else:
            print("   (来源: config.py配置文件)")
            
        print(f"🌐 目标网址: {target_url}")
        if args.target_url:
            print("   (来源: 命令行参数)")
        else:
            print("   (来源: config.py配置文件)")
            
        print(f"🔌 插件备份: {'禁用' if args.no_extensions else '启用'}")
        print()
            
        if not username:
            print("❌ 错误：未设置用户名")
            print("请使用 -u 参数指定用户名，或在config.py中配置 USERNAME")
            input("按回车键退出...")
            exit(1)
        
        success = legacy_login_with_extensions(username, password, target_url, backup_extensions)
        if success:
            print("\n🎉 登录完成！")
            print(f"💾 网站登录状态文件: {username}_auth_state.json")
            if backup_extensions:
                print(f"📦 扩展程序备份文件: {username}_extensions_backup.zip")
            print("\n💡 现在您可以：")
            print("   1. 在其他电脑上运行 open_with_extensions.py 获得完整环境")
            print("   2. 或运行普通的 open.py 仅获得网站登录状态")
        else:
            print("\n❌ 登录失败，请检查配置并重试")
            
    except KeyboardInterrupt:
        print("\n🚫 用户取消操作")
    except Exception as e:
        print(f"\n💥 程序遇到未知错误: {e}")
        import traceback
        traceback.print_exc()
    finally:
        input("\n按回车键退出...")